from sklearn import datasets
import torch

iris = datasets.load_iris()
x, y = torch.tensor(iris.data, dtype=torch.float32), torch.tensor(iris.target, dtype=torch.long)

INPUT_DIM = 4  # входные параметры
OUTPUT_DIM = 3  # ответы нейронной сети

H_DIM = 5  # Количество нейронов в слое

EPOCH = 100 # Количество эпох
Learning_rate = 0.5 # скорость обучения

two_layer_net = torch.nn.Sequential( # структура нейросети
    torch.nn.Linear(INPUT_DIM, H_DIM),
    torch.nn.ReLU(),
    torch.nn.Linear(H_DIM, OUTPUT_DIM),
    torch.nn.Softmax(dim=-1)
)

loss_func = torch.nn.CrossEntropyLoss() # ошибка
optim = torch.optim.SGD(two_layer_net.parameters(), lr=Learning_rate)

history = []

for epoch in range(EPOCH):

    #Forward
    y_pred = two_layer_net(x)

    #Loss
    loss = loss_func(y_pred, y)

    #Backward
    loss.backward()

    #Learning
    optim.step()

    optim.zero_grad()

    history.append(loss.item())

import matplotlib.pyplot as plt
plt.plot(history,label='loss')
plt.show()

x_new = torch.tensor([7.1, 3.5 , 4.7, 1.3], dtype=torch.float32)
class_iris = ['setosa', 'versicolor', 'virginica']
print(class_iris[torch.argmax(two_layer_net(x_new))])

def calc_accuracy():
    correct = 0
    for i in range(len(y)):
        z = two_layer_net(x[i])
        y_pred = torch.argmax(z)
        if y_pred == y[i]:
            correct += 1
    acc = correct / len(y)
    return acc


accuracy = calc_accuracy()
print("Accuracy: ", accuracy)
